{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# rStar-Math: Statistics and Probability Examples\n",
    "\n",
    "This notebook demonstrates statistical analysis and probability problems with visualizations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import stats\n",
    "from src.core.mcts import MCTS\n",
    "from src.core.ppm import ProcessPreferenceModel\n",
    "from src.models.model_interface import ModelFactory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "# Initialize components\n",
    "mcts = MCTS.from_config_file('config/default.json')\n",
    "ppm = ProcessPreferenceModel.from_config_file('config/default.json')\n",
    "model = ModelFactory.create_model('openai', os.getenv('OPENAI_API_KEY'), 'config/default.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Descriptive Statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "stats_problems = [\n",
    "    \"Find the mean, median, and mode of [2, 3, 3, 4, 4, 4, 5, 5, 6]\",\n",
    "    \"Calculate the standard deviation of [10, 12, 15, 18, 20]\",\n",
    "    \"Find the quartiles and IQR of [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]\"\n",
    "]\n",
    "\n",
    "def visualize_distribution(data: list):\n",
    "    \"\"\"Visualize data distribution.\"\"\"\n",
    "    plt.figure(figsize=(12, 4))\n",
    "    \n",
    "    # Histogram\n",
    "    plt.subplot(131)\n",
    "    plt.hist(data, bins='auto', alpha=0.7)\n",
    "    plt.title('Histogram')\n",
    "    \n",
    "    # Box plot\n",
    "    plt.subplot(132)\n",
    "    plt.boxplot(data)\n",
    "    plt.title('Box Plot')\n",
    "    \n",
    "    # KDE plot\n",
    "    plt.subplot(133)\n",
    "    sns.kdeplot(data=data)\n",
    "    plt.title('Density Plot')\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "for problem in stats_problems:\n",
    "    print(f\"Problem: {problem}\\n\")\n",
    "    action, trajectory = mcts.search(problem)\n",
    "    \n",
    "    print(\"Solution Steps:\")\n",
    "    for step in trajectory:\n",
    "        confidence = ppm.evaluate_step(step['state'], model)\n",
    "        print(f\"- {step['state']}\")\n",
    "        print(f\"  Confidence: {confidence:.2f}\\n\")\n",
    "        \n",
    "    # Visualize data if present\n",
    "    if '[' in problem:\n",
    "        data = eval(problem[problem.find('['):problem.find(']')+1])\n",
    "        visualize_distribution(data)\n",
    "    print(\"-\" * 50 + \"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Probability Distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "def plot_probability_distribution(dist_type: str, params: dict):\n",
    "    \"\"\"Plot various probability distributions.\"\"\"\n",
    "    plt.figure(figsize=(10, 6))\n",
    "    \n",
    "    if dist_type == 'normal':\n",
    "        x = np.linspace(params['mu'] - 4*params['sigma'],\n",
    "                       params['mu'] + 4*params['sigma'], 100)\n",
    "        y = stats.norm.pdf(x, params['mu'], params['sigma'])\n",
    "        plt.plot(x, y, label=f'μ={params[\"mu\"]}, σ={params[\"sigma\"]}')\n",
    "        plt.title('Normal Distribution')\n",
    "        \n",
    "    elif dist_type == 'binomial':\n",
    "        x = np.arange(0, params['n']+1)\n",
    "        y = stats.binom.pmf(x, params['n'], params['p'])\n",
    "        plt.bar(x, y, alpha=0.8, label=f'n={params[\"n\"]}, p={params[\"p\"]}')\n",
    "        plt.title('Binomial Distribution')\n",
    "        \n",
    "    elif dist_type == 'poisson':\n",
    "        x = np.arange(0, params['lambda']*3)\n",
    "        y = stats.poisson.pmf(x, params['lambda'])\n",
    "        plt.bar(x, y, alpha=0.8, label=f'λ={params[\"lambda\"]}')\n",
    "        plt.title('Poisson Distribution')\n",
    "    \n",
    "    plt.grid(True, alpha=0.3)\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "\n",
    "# Example distributions\n",
    "distributions = [\n",
    "    ('normal', {'mu': 0, 'sigma': 1}),\n",
    "    ('normal', {'mu': 0, 'sigma': 2}),\n",
    "    ('binomial', {'n': 20, 'p': 0.3}),\n",
    "    ('poisson', {'lambda': 3})\n",
    "]\n",
    "\n",
    "for dist_type, params in distributions:\n",
    "    plot_probability_distribution(dist_type, params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Hypothesis Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "def visualize_hypothesis_test(sample1: np.ndarray, sample2: np.ndarray, test_type: str):\n",
    "    \"\"\"Visualize hypothesis test results.\"\"\"\n",
    "    plt.figure(figsize=(12, 5))\n",
    "    \n",
    "    # Data distribution\n",
    "    plt.subplot(121)\n",
    "    plt.boxplot([sample1, sample2], labels=['Sample 1', 'Sample 2'])\n",
    "    plt.title('Sample Distributions')\n",
    "    \n",
    "    # Test results\n",
    "    if test_type == 't-test':\n",
    "        t_stat, p_val = stats.ttest_ind(sample1, sample2)\n",
    "        test_name = \"Student's t-test\"\n",
    "    elif test_type == 'wilcoxon':\n",
    "        t_stat, p_val = stats.wilcoxon(sample1, sample2)\n",
    "        test_name = \"Wilcoxon signed-rank test\"\n",
    "    \n",
    "    plt.subplot(122)\n",
    "    plt.text(0.5, 0.5,\n",
    "             f\"Test: {test_name}\\n\" +\n",
    "             f\"Statistic: {t_stat:.4f}\\n\" +\n",
    "             f\"p-value: {p_val:.4f}\\n\" +\n",
    "             f\"Significant: {p_val < 0.05}\",\n",
    "             ha='center', va='center')\n",
    "    plt.axis('off')\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# Example hypothesis tests\n",
    "np.random.seed(42)\n",
    "sample1 = np.random.normal(0, 1, 100)\n",
    "sample2 = np.random.normal(0.5, 1, 100)\n",
    "\n",
    "print(\"Comparing two samples with different means:\")\n",
    "visualize_hypothesis_test(sample1, sample2, 't-test')\n",
    "\n",
    "print(\"\\nNon-parametric test:\")\n",
    "visualize_hypothesis_test(sample1, sample2, 'wilcoxon')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
